# text2sql_evaluator_qwen_openai.py
import os
from typing import List, Dict, Any
import sqlite3
from langchain_community.utilities import SQLDatabase
from openai import OpenAI
import re

class Text2SQLEvaluator:
    def __init__(self, db_uri: str, base_url: str = "http://localhost:8000/v1", api_key: str = "EMPTY", model_name: str = "Qwen2.5-7B-Instruct", temperature: float = 0.0):
        """
        初始化评估器，通过 OpenAI 兼容 API 调用 Qwen2.5-7B 模型
        :param db_uri: 数据库连接 URI，例如 "sqlite:///sample.db"
        :param base_url: OpenAI API 的基础 URL（例如本地 vLLM 服务）
        :param api_key: API 密钥（本地部署通常为 "EMPTY"）
        :param model_name: 模型名称（需与部署的服务中模型名称一致）
        :param temperature: 生成温度
        """
        # 连接数据库
        self.db = SQLDatabase.from_uri(db_uri)
        
        # 初始化 OpenAI 客户端（兼容本地部署的 Qwen）
        self.client = OpenAI(
            base_url=base_url,
            api_key=api_key
        )
        
        self.model_name = model_name
        self.temperature = temperature

        # 系统提示词（System Prompt）
        self.system_prompt = """你是一个专业的 SQL 生成助手。你的任务是根据用户的自然语言问题和给定的数据库表结构，生成准确、简洁的 SQL 查询语句。
要求：
1. 只生成 SQL 语句，不要包含任何解释、注释或其他文本。
2. SQL 语句必须能在 SQLite 数据库中正确执行。
3. 如果问题无法回答或信息不足，请返回 "SELECT '无法生成SQL' AS error;"。
4. 请确保 SQL 语法正确。

当前数据库包含以下表：
{table_info}

请根据以上信息生成 SQL。
"""

    def _format_messages(self, question: str) -> list:
        """
        格式化为 OpenAI ChatCompletion 所需的消息格式
        """
        table_info = self.db.get_table_info()
        system_content = self.system_prompt.format(table_info=table_info)

        messages = [
            {"role": "system", "content": system_content},
            {"role": "user", "content": question}
        ]
        return messages

    def generate_sql(self, question: str) -> str:
        """
        调用 Qwen2.5 模型生成 SQL
        """
        try:
            messages = self._format_messages(question)

            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                temperature=self.temperature,
                max_tokens=512,
                top_p=0.95,
                stream=False
            )

            generated_text = response.choices[0].message.content.strip()

            # 清理输出：移除 Markdown 代码块
            if generated_text.startswith("```sql"):
                generated_text = generated_text[6:].strip()
            if generated_text.startswith("```"):
                generated_text = generated_text[3:].strip()
            if generated_text.endswith("```"):
                generated_text = generated_text[:-3].strip()

            # 截取到第一个分号（防止多语句或附加文本）
            semicolon_index = generated_text.find(';')
            if semicolon_index != -1:
                generated_sql = generated_text[:semicolon_index + 1].strip()
            else:
                generated_sql = generated_text.strip()

            return generated_sql

        except Exception as e:
            print(f"❌ 生成 SQL 时出错: {e}")
            return "SELECT 'ERROR' AS error;"

    def execute_sql(self, sql: str) -> str:
        """
        执行 SQL 并返回结果字符串
        """
        try:
            result = self.db.run(sql)
            return str(result).strip()
        except Exception as e:
            return f"EXECUTION_ERROR: {str(e)}"

    def evaluate_execution_accuracy(self, generated_sql: str, gold_sql: str) -> bool:
        """
        通过执行结果判断是否正确
        """
        try:
            gen_result = self.execute_sql(generated_sql)
            gold_result = self.execute_sql(gold_sql)
            return gen_result == gold_result
        except Exception as e:
            print(f"评估时出错: {e}")
            return False

    def evaluate(self, test_dataset: List[Dict[str, str]]) -> Dict[str, Any]:
        """
        评估整个测试集
        """
        total = len(test_dataset)
        correct = 0
        results = []

        print(f"🧪 开始评估 {total} 个问题...")

        for i, item in enumerate(test_dataset):
            question = item["question"]
            gold_sql = item["gold_sql"]

            print(f"\n--- [{i+1}/{total}] 问题: {question} ---")

            # 生成 SQL
            generated_sql = self.generate_sql(question)
            print(f"🧠 生成 SQL: {generated_sql}")

            # 评估执行准确性
            is_correct = self.evaluate_execution_accuracy(generated_sql, gold_sql)
            if is_correct:
                correct += 1
                print("✅ 执行结果匹配")
            else:
                print("❌ 执行结果不匹配")

            results.append({
                "question": question,
                "generated_sql": generated_sql,
                "gold_sql": gold_sql,
                "is_correct": is_correct
            })

        execution_accuracy = correct / total if total > 0 else 0.0

        return {
            "execution_accuracy": execution_accuracy,
            "total_questions": total,
            "correct_count": correct,
            "detailed_results": results
        }


# 主程序
if __name__ == "__main__":
    from test_dataset import TEST_DATASET

    # ========================
    # 🔧 配置您的 Qwen2.5 API 地址
    # ========================
    # 示例 1：本地 vLLM 默认端口
    BASE_URL = "http://10.10.202.242:9000/v1"
    API_KEY = "EMPTY"
    MODEL_NAME = "Qwen2.5-14B-Instruct-GPTQ-Int4"  # 必须与您启动服务时指定的模型名称一致

    # 示例 2：如果您使用阿里云百炼或其他平台，请替换为对应地址和密钥
    # BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
    # API_KEY = os.getenv("DASHSCOPE_API_KEY")
    # MODEL_NAME = "qwen2.5-7b-instruct"

    # 初始化评估器
    evaluator = Text2SQLEvaluator(
        db_uri="sqlite:///sample.db",
        base_url=BASE_URL,
        api_key=API_KEY,
        model_name=MODEL_NAME,
        temperature=0.0
    )

    # 运行评估
    print("🚀 开始评估 Text2SQL 模型 (Qwen2.5-7B via OpenAI API)...")
    results = evaluator.evaluate(TEST_DATASET)

    # 输出报告
    print("\n" + "="*60)
    print("📊 最终评估报告")
    print("="*60)
    print(f"总问题数: {results['total_questions']}")
    print(f"正确执行数: {results['correct_count']}")
    print(f"✅ 执行准确率: {results['execution_accuracy']:.2%}")

    print("\n📝 详细结果:")
    for i, res in enumerate(results['detailed_results'], 1):
        mark = "✅" if res['is_correct'] else "❌"
        print(f"{i:2d}. {mark} {res['question']}")